# 1、数据预处理
import numpy as np
import pandas as pd

data = pd.read_csv('zgpa_train.csv')
# print(笔记.md.head())
price = data.loc[:, 'close']
# print(price.head())
# 归一化处理
price_norm = price / max(price)
# print(price_norm.head())
# 可视化操作
from matplotlib import pyplot as plt

fig1 = plt.figure(figsize=(8, 5))
plt.plot(price)
plt.title('收盘价格', fontproperties='SimHei', fontsize=20)
plt.xlabel('时间', fontproperties='SimHei', fontsize=20)
plt.ylabel('价格', fontproperties='SimHei', fontsize=20)


# plt.show()


# 提取x,y
def extract_data(data, time_step):
    x = []
    y = []
    # 0,1,2,3,4,5,6,7,8,9 10个样本，time_step=8
    for i in range(len(data) - time_step):
        x.append([a for a in data[i:i + time_step]])
        y.append(data[i + time_step])
    x = np.array(x)
    x = x.reshape(x.shape[0], x.shape[1], 1)
    return x, np.array(y)


time_step = 8
x, y = extract_data(price_norm, time_step)
# print(x)


# 2、建立模型进行预测
from keras.models import Sequential
from keras.layers import Dense, SimpleRNN

model = Sequential()
model.add(SimpleRNN(units=5, input_shape=(time_step, 1), activation='relu'))
model.add(Dense(units=1, activation='linear'))
# 配置模型
model.compile(optimizer='adam', loss='mean_squared_error')
model.summary()
# 模型训练
model.fit(x,y,batch_size=30,epochs=200)
# 可视化预测效果
y_train_predict = model.predict(x)*max(price)
y_train = y*max(price)
# print(y_train_predict,y_train)
fig2 = plt.figure(figsize=(8, 5))
plt.plot(y_train,label='y_train')
plt.plot(y_train_predict,label='y_train_predict')
plt.title('收盘价格', fontproperties='SimHei', fontsize=20)
plt.xlabel('时间', fontproperties='SimHei', fontsize=20)
plt.ylabel('价格', fontproperties='SimHei', fontsize=20)
plt.legend()
# plt.show()

# 对测试数据进行预测
data_test = pd.read_csv('zgpa_test.csv')
price_test = data_test.loc[:,'close']
price_test_norm = price_test/max(price) # 归一化
x_test_norm,y_test_norm = extract_data(price_test_norm,time_step)

y_test_predict = model.predict(x_test_norm)*max(price)
y_test = y_test_norm*max(price)

fig3 = plt.figure(figsize=(8, 5))
plt.plot(y_test,label='y_test')
plt.plot(y_test_predict,label='y_test_predict')
plt.title('收盘价格', fontproperties='SimHei', fontsize=20)
plt.xlabel('时间', fontproperties='SimHei', fontsize=20)
plt.ylabel('价格', fontproperties='SimHei', fontsize=20)
plt.legend()
# plt.show()

# 存储数据
result_y_test = np.array(y_test).reshape(-1,1)
result_y_test_predict = y_test_predict
result = np.concatenate((result_y_test,result_y_test_predict),axis=1)
result = pd.DataFrame(result,columns=['real-price_test','predict_price_test'])
result.to_csv('zgpa_predict_test.csv')
